
#include "preprocess_kernel.cuh"

namespace CUDAKernel
{
    Norm Norm::mean_std(const float mean[3], const float std[3], float alpha, ChannelType channel_type)
    {

        Norm out;
        out.type = NormType::MeanStd;
        out.alpha = alpha;
        out.channel_type = channel_type;
        memcpy(out.mean, mean, sizeof(out.mean));
        memcpy(out.std, std, sizeof(out.std));
        return out;
    }

    Norm Norm::alpha_beta(float alpha, float beta, ChannelType channel_type)
    {

        Norm out;
        out.type = NormType::AlphaBeta;
        out.alpha = alpha;
        out.beta = beta;
        out.channel_type = channel_type;
        return out;
    }

    Norm Norm::None()
    {
        return Norm();
    }

    __global__ void resize_bilinear_and_normalize_kernel(
        uint8_t *src, int src_line_size, int src_width, int src_height, float *dst, int dst_width, int dst_height,
        float sx, float sy, Norm norm, int threads)
    {
        int position = blockDim.x * blockIdx.x + threadIdx.x;
        if (position >= threads)
            return;

        int dx = position % dst_width;
        int dy = position / dst_width;
        float src_x = (dx + 0.5f) * sx - 0.5f;
        float src_y = (dy + 0.5f) * sy - 0.5f;
        float c0, c1, c2;

        int y_low = floorf(src_y);
        int x_low = floorf(src_x);
        int y_high = clipf(y_low + 1, 0, src_height - 1);
        int x_high = clipf(x_low + 1, 0, src_width - 1);
        y_low = clipf(y_low, 0, src_height - 1);
        x_low = clipf(x_low, 0, src_width - 1);

        int ly = rint((src_y - y_low) * INTER_RESIZE_COEF_SCALE);
        int lx = rint((src_x - x_low) * INTER_RESIZE_COEF_SCALE);
        int hy = INTER_RESIZE_COEF_SCALE - ly;
        int hx = INTER_RESIZE_COEF_SCALE - lx;
        int w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
        float *pdst = dst + dy * dst_width + dx * 3;
        uint8_t *v1 = src + y_low * src_line_size + x_low * 3;
        uint8_t *v2 = src + y_low * src_line_size + x_high * 3;
        uint8_t *v3 = src + y_high * src_line_size + x_low * 3;
        uint8_t *v4 = src + y_high * src_line_size + x_high * 3;

        c0 = resize_cast(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]);
        c1 = resize_cast(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]);
        c2 = resize_cast(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]);

        if (norm.channel_type == ChannelType::Invert)
        {
            float t = c2;
            c2 = c0;
            c0 = t;
        }

        if (norm.type == NormType::MeanStd)
        {
            c0 = (c0 * norm.alpha - norm.mean[0]) / norm.std[0];
            c1 = (c1 * norm.alpha - norm.mean[1]) / norm.std[1];
            c2 = (c2 * norm.alpha - norm.mean[2]) / norm.std[2];
        }
        else if (norm.type == NormType::AlphaBeta)
        {
            c0 = c0 * norm.alpha + norm.beta;
            c1 = c1 * norm.alpha + norm.beta;
            c2 = c2 * norm.alpha + norm.beta;
        }

        int area = dst_width * dst_height;
        float *pdst_c0 = dst + dy * dst_width + dx;
        float *pdst_c1 = pdst_c0 + area;
        float *pdst_c2 = pdst_c1 + area;
        *pdst_c0 = c0;
        *pdst_c1 = c1;
        *pdst_c2 = c2;
    }

    __global__ void resize_bilinear_kernel(
        uint8_t *src, int src_line_size, int src_width, int src_height, float *dst, int dst_width, int dst_height,
        float sx, float sy, int threads)
    {
        int position = threadIdx.x + blockDim.x * blockIdx.x;
        if (position >= threads)
            return;

        int dx = position % dst_width;
        int dy = position / dst_width;
        // (dx + 0.5)*sx - 0.5 的主要原因是使得中心更加对齐，下采样不明显，但是上采样时就比较明显
        float src_x = (dx + 0.5f) * sx - 0.5f; // sx = src_width / (float)dst_width
        float src_y = (dy + 0.5f) * sy - 0.5f;
        float c0, c1, c2;

        // 寻找目标点dst对应src中相邻的四个点
        int y_low = floorf(src_y);
        int x_low = floorf(src_x);
        int y_high = clipf(y_low + 1, 0, src_height - 1);
        int x_high = clipf(x_low + 1, 0, src_width - 1);
        y_low = clipf(y_low, 0, src_height - 1);
        x_low = clipf(x_low, 0, src_width - 1);

        // 求解四点权重，但都放大了 INTER_RESIZE_COEF_SCALE倍，目的是在四舍五入的时候减少损失
        int ly = rint((src_y - y_low) * INTER_RESIZE_COEF_SCALE);
        int lx = rint((src_x - x_low) * INTER_RESIZE_COEF_SCALE);
        int hy = INTER_RESIZE_COEF_SCALE - ly;
        int hx = INTER_RESIZE_COEF_SCALE - lx;
        int w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

        uint8_t *v1 = src + y_low * src_line_size + x_low * 3;
        uint8_t *v2 = src + y_low * src_line_size + x_high * 3;
        uint8_t *v3 = src + y_high * src_line_size + x_low * 3;
        uint8_t *v4 = src + y_high * src_line_size + x_high * 3;

        // resize_cast是为了把前面放大的INTER_RESIZE_COEF_SCALE倍 恢复原样，这样损失就超小
        c0 = resize_cast(w1 * v1[0] + w2 * v2[0] + w3 * v3[0] + w4 * v4[0]);
        c1 = resize_cast(w1 * v1[1] + w2 * v2[1] + w3 * v3[1] + w4 * v4[1]);
        c2 = resize_cast(w1 * v1[2] + w2 * v2[2] + w3 * v3[2] + w4 * v4[2]);

        int area = dst_width * dst_height;
        float *pdst_c0 = dst + dy * dst_width + dx;
        float *pdst_c1 = pdst_c0 + area;
        float *pdst_c2 = pdst_c1 + area;
        *pdst_c0 = c0; // 指针解应用，赋值操作
        *pdst_c1 = c1; // 指针解应用，赋值操作
        *pdst_c2 = c2; // 指针解应用，赋值操作
    }

    void resize_bilinear(uint8_t *src, int src_line_size, int src_width, int src_height,
                         float *dst, int dst_width, int dst_height, cudaStream_t stream)
    {

        // 启用 w*h个线程，每个线程计算b,g,r三个通道
        int threads = dst_width * dst_height;
        auto grid = grid_dims(threads);
        auto block = block_dims(threads);

        resize_bilinear_kernel<<<grid, block, 0, stream>>>(
            src, src_line_size,
            src_width, src_height, dst,
            dst_width, dst_height, src_width / (float)dst_width, src_height / (float)dst_height, threads);

        // auto norm = CUDAKernel::Norm::None(); // 一般在最外面定义，这里只是做个简单的演示
        // resize_bilinear_and_normalize_kernel<<<grid, block, 0, stream>>>(
        //     src, src_line_size,
        //     src_width, src_height, dst,
        //     dst_width, dst_height, src_width / (float)dst_width, src_height / (float)dst_height, norm, threads);
    }

};
